package com.tuniu.agents.common.logger;

import com.tuniu.agents.advisor.ChatEnvironmentMemory;
import com.tuniu.agents.common.AgentsConstants;
import com.tuniu.agents.message.MessageLogger;
import com.tuniu.agents.message.MessageSender;
import com.tuniu.agents.message.ToolLoggerInfo;
import lombok.extern.slf4j.Slf4j;
import org.apache.commons.lang3.StringUtils;
import org.aspectj.lang.ProceedingJoinPoint;
import org.aspectj.lang.annotation.Around;
import org.aspectj.lang.annotation.Aspect;
import org.aspectj.lang.reflect.MethodSignature;
import org.springframework.ai.chat.model.ToolContext;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.context.ApplicationContext;
import org.springframework.core.annotation.Order;
import org.springframework.stereotype.Component;

import java.lang.reflect.Method;
import java.util.Map;

@Aspect
@Component
@Order(1)
@Slf4j
public class ToolLoggerAspect extends LoggerAspect{

    @Autowired
    private ChatEnvironmentMemory chatEnvironmentMemory;

    public ToolLoggerAspect(ApplicationContext applicationContext, MessageSender messageSender) {
        super(applicationContext, messageSender);
    }

    @Around("@annotation(com.tuniu.agents.message.ToolLoggerInfo) || @within(com.tuniu.agents.message.ToolLoggerInfo)")
    public Object aroundLogExecution(ProceedingJoinPoint joinPoint) throws Throwable {
        // Get method signature
        MethodSignature signature = (MethodSignature) joinPoint.getSignature();
        Method method = signature.getMethod();

        String conversationId = this.getConversationId(joinPoint.getArgs());

        if(StringUtils.isNotBlank(conversationId)) {
            // Get all Bean names
            try {
                String[] beanNames = super.applicationContext.getBeanNamesForType(joinPoint.getTarget().getClass());
                if(beanNames != null && beanNames.length > 0) {
                    // Set tool ID
                    this.chatEnvironmentMemory.add(conversationId, Map.of(AgentsConstants.AgentsConversationEnvironmentKey.LAST_SELECTION_TOOL_ID, beanNames[0]));
                } else {
                    this.chatEnvironmentMemory.add(conversationId, Map.of(AgentsConstants.AgentsConversationEnvironmentKey.LAST_SELECTION_TOOL_ID, ""));
                }
            }catch (Exception e) {
                log.error("Error getting tool name", e);
            }
        }

        // First try to get the annotation from the method
        ToolLoggerInfo logAnnotation = method.getAnnotation(ToolLoggerInfo.class);

        // If there is no annotation on the method, try to get it from the class
        if (logAnnotation == null) {
            logAnnotation = joinPoint.getTarget().getClass().getAnnotation(ToolLoggerInfo.class);
        }

        MessageLogger logger = logAnnotation == null ? null : MessageLogger.fromToolLoggerInfo(logAnnotation);

        Object result = getResult(joinPoint, logger);

        return result;
    }

    @Override
    public String getConversationId(Object[] args) {
        Object conversationId = ((ToolContext) args[1]).getContext().get(AgentsConstants.AgentsConversationEnvironmentKey.CONVERSATION_ID);
        return conversationId == null ? null : String.valueOf(conversationId);
    }
}
